import random
import numpy as np
from keras.utils import np_utils

from config import Config


class GarbageDataSet:
    def load_data(self, filename: str, testRate: float = 0.2):
        imgdata = np.load(filename)
        # 分组 16个一组
        group_h = [imgdata[0:16000][16 * i:16 * (i + 1)] for i in range(1000)]
        group_k = [imgdata[16000:32000][16 * i:16 * (i + 1)] for i in range(1000)]
        group_o = [imgdata[32000:48000][16 * i:16 * (i + 1)] for i in range(1000)]
        group_r = [imgdata[48000:64000][16 * i:16 * (i + 1)] for i in range(1000)]
        # 打乱顺序
        # print(np.array(group_h).shape)
        random.shuffle(group_h)
        random.shuffle(group_k)
        random.shuffle(group_o)
        random.shuffle(group_r)
        # 分测试集和训练集 把分组再变回一个np数组
        split_val = (int(16000 * (1 - testRate)) // 16)
        x_train = np.array(
            [item for obj in [group_h[:split_val],
                              group_k[:split_val],
                              group_o[:split_val],
                              group_r[:split_val]] for group in obj for item in group]
        )
        x_test = np.array(
            [item for obj in [group_h[split_val:],
                              group_k[split_val:],
                              group_o[split_val:],
                              group_r[split_val:]] for group in obj for item in group]
        )
        # map
        # 1 0 0 0   h 0
        # 0 1 0 0   k 1
        # 0 0 1 0   o 2
        # 0 0 0 1   r 3
        y_train = np_utils.to_categorical(np.array(self._get_y(len(x_train))), Config.NUM_CLASSES)
        y_test = np_utils.to_categorical(np.array(self._get_y(len(x_test))), Config.NUM_CLASSES)
        return x_train, y_train, x_test, y_test

    def _get_y(self, length):
        res = []
        for i in range(Config.NUM_CLASSES):
            res += [i] * (length // Config.NUM_CLASSES)
        return res

    def load_data_from_npy(self):
        x_train = np.load(f"x_train_float{Config.IMAGE_HEIGHT}.npy")
        x_test = np.load(f"x_test_float{Config.IMAGE_HEIGHT}.npy")
        y_train = np.load(f"y_train{Config.IMAGE_HEIGHT}.npy")
        y_test = np.load(f"y_test{Config.IMAGE_HEIGHT}.npy")
        return x_train, y_train, x_test, y_test

    def normalization(self, x_train, y_train, x_test, y_test, save: bool = True):  # 归一化
        x_train = x_train.astype('float32')
        x_test = x_test.astype('float32')
        x_train /= 255
        x_test /= 255
        if save:
            np.save(f"x_train_float{len(x_train[1])}.npy", x_train)
            np.save(f"x_test_float{len(x_train[1])}.npy", x_test)
            np.save(f"y_train{len(x_train[1])}.npy", y_train)
            np.save(f"y_test{len(x_train[1])}.npy", y_test)


if __name__ == "__main__":
    g = GarbageDataSet()
    a, b, c, d = g.load_data("data/data32.npy")
    print(a.shape)
    print(b.shape)
    print(c.shape)
    print(d.shape)
